from optimum.exporters.onnx import main_export
from transformers import AutoConfig
import os
os.environ["HTTP_PROXY"] = "http://127.0.0.1:7897"
os.environ["HTTPS_PROXY"] = "http://127.0.0.1:7897"
def export_chatglm_comprehensive():
    model_id = "zai-org/chatglm3-6b"
    output_dir = "./chatglm3-onnx-final"

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    try:
        # 尝试多种配置
        main_export(
            model_name_or_path=model_id,
            output=output_dir,
            task="text-generation",  # 最通用的任务
            trust_remote_code=True,  # 必须的
            opset=14,  # 合适的opset版本
            device="cpu",  # 在CPU上导出更稳定
            do_constant_folding=True,  # 常量折叠优化
        )
        print(f"ChatGLM 模型成功导出到: {output_dir}")

    except Exception as e:
        print(f"导出失败: {e}")
        print("尝试备用方案...")

        # 备用方案：使用默认配置
        try:
            main_export(
                model_name_or_path=model_id,
                output=output_dir + "_fallback",
                trust_remote_code=True,
                opset=14,
                device="cpu",
            )
            print("备用方案导出成功！")
        except Exception as e2:
            print(f"所有导出尝试都失败: {e2}")


if __name__ == "__main__":
    export_chatglm_comprehensive()